#!/usr/bin/env python
# coding: utf-8

# ======= dynamic modeling =====
import scvelo as scv

# prepare the data ----------
adata = scv.read('/home/gjsx/learn/scvelo/data/Pancreas/endocrinogenesis_day15.h5ad', cache=True)

scv.pp.filter_and_normalize(adata, min_shared_counts=20, n_top_genes=2000)
scv.pp.moments(adata, n_pcs=30)

# dynamical model -------
# aims to learn the unspliced/spliced phase trajectory for each gene.
scv.tl.recover_dynamics(adata, n_jobs=16)
# recover_dynamics can take a while
# ~10min for 1 core, 1min 13s for 16 cores

scv.tl.velocity(adata, mode='dynamical')
scv.tl.velocity_graph(adata)

# better save product on disk
adata.write('/home/gjsx/learn/scvelo/data/panc_dynamic.h5ad', compression='gzip')
adata = scv.read('/home/gjsx/learn/scvelo/data/panc_dynamic.h5ad')

scv.pl.velocity_embedding_stream(adata, basis='umap')

# Kinetic rate parameter ----------
# rates of RNA transcription, splicing and degradation are estimated
# extract var (gene) info
df = adata.var
# filter genes with high fit likelihood and with velocity
df = df[(df['fit_likelihood'] > .1) & df['velocity_genes'] == True]

# fit_alpha = transcription rate
# fit_beta * fit_scaling = splicing rate
# fit_gamma = degradation rate
kwargs = dict(xscale='log', fontsize=16)
with scv.GridSpec(ncols=3) as pl:
    pl.hist(df['fit_alpha'], xlabel='transcription rate', **kwargs)
    pl.hist(df['fit_beta'] * df['fit_scaling'], xlabel='splicing rate', xticks=[.1, .4, 1], **kwargs)
    pl.hist(df['fit_gamma'], xlabel='degradation rate', xticks=[.1, .4, 1], **kwargs)

# select all column named with fit
scv.get_df(adata, 'fit*', dropna=True).head()

# Latent time ---------
# approximates the real time experienced by cells as they differentiate
scv.tl.latent_time(adata)
scv.pl.scatter(adata, color='latent_time', color_map='gnuplot', size=80)

# filter 300 best fit genes
top_genes = adata.var['fit_likelihood'].sort_values(ascending=False).index[:300]
# plot genes along time axis
scv.pl.heatmap(adata, var_names=top_genes, sortby='latent_time', col_color='clusters', n_convolve=100)

# Top-likelihood genes ----------
# Driver genes display pronounced dynamic behavior
top_genes = adata.var['fit_likelihood'].sort_values(ascending=False).index
# plot 15 best fit genes
scv.pl.scatter(adata, basis=top_genes[:15], ncols=5, frameon=False)

var_names = ['Actn4', 'Ppp3ca', 'Cpe', 'Nnat']
scv.pl.scatter(adata, x='latent_time', y='Cpe')

# Cluster-specific top-likelihood genes -----------
# driver gene likelihoods can be computed for each cluster of cells
scv.tl.rank_dynamical_genes(adata, groupby='clusters')
df = scv.get_df(adata, 'rank_dynamical_genes/names')
df.head(5)

for cluster in ['Ductal', 'Ngn3 high EP', 'Pre-endocrine', 'Beta']:
    scv.pl.scatter(adata, df[cluster][:5], ylabel=cluster, frameon=False)
